package main

import (
	"bufio"
	"fmt"
	"log"
	"os"
	"strconv"
	"strings"

	"github.com/spf13/cobra"
	"golang.org/x/crypto/ssh"
	"github.com/username/go-demo/internal/app"
)

// 添加日志器实现
func GetLogger() *log.Logger {
	return log.New(os.Stderr, "[BASTION] ", log.LstdFlags)
}

func main() {
	var rootCmd = &cobra.Command{
		Use:   "bastion",
		Short: "命令行堡垒机管理系统",
	}

	rootCmd.AddCommand(
		newAddHostCommand(),
		newListHostsCommand(),
		newConnectCommand(),
	)

	if err := rootCmd.Execute(); err != nil {
		app.GetLogger().Fatal(err)
	}
}

func newAddHostCommand() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "add-host",
		Short: "添加远程主机",
		Run: func(cmd *cobra.Command, args []string) {
			storage := app.NewStorageService()

			hosts, _ := storage.LoadHosts()
			newHost := app.Host{
				IP:       prompt("请输入IP地址: "),
				Username: prompt("请输入用户名: "),
				Password: prompt("请输入密码: "),
				Port:     mustAtoi(prompt("请输入端口: ")),
			}

			hosts = append(hosts, newHost)
			if err := storage.SaveHosts(hosts); err != nil {
				app.GetLogger().Fatal("保存失败:", err)
			}
			fmt.Println("\n主机添加成功!")
		},
	}
	return cmd
}

func newListHostsCommand() *cobra.Command {
	return &cobra.Command{
		Use:   "list-hosts",
		Short: "列出所有已存主机",
		Run: func(cmd *cobra.Command, args []string) {
			storage := app.NewStorageService()
			hosts, err := storage.LoadHosts()
			if err != nil {
				app.GetLogger().Fatal("加载失败:", err)
			}

			fmt.Println("\n已保存主机列表:")
			for i, h := range hosts {
				fmt.Printf("%d) %s@%s:%d\n", i+1, h.Username, h.IP, h.Port)
			}
		},
	}
}

func newConnectCommand() *cobra.Command {
	return &cobra.Command{
		Use:   "connect",
		Short: "连接指定主机",
		Run: func(cmd *cobra.Command, args []string) {
			storage := app.NewStorageService()
			hosts, err := storage.LoadHosts()
			if err != nil {
				app.GetLogger().Fatal("加载失败:", err)
			}

			// 显示主机列表
			fmt.Println("\n可用主机列表:")
			for i, h := range hosts {
				fmt.Printf("%d) %s@%s:%d\n", i+1, h.Username, h.IP, h.Port)
			}

			// 选择主机
			selected := mustAtoi(prompt("\n请输入要连接的主机编号: ")) - 1
			if selected < 0 || selected >= len(hosts) {
				app.GetLogger().Fatal("无效的选择")
			}

			host := hosts[selected]
			config := &ssh.ClientConfig{
				User: host.Username,
				Auth: []ssh.AuthMethod{
					ssh.Password(host.Password),
				},
				HostKeyCallback: ssh.InsecureIgnoreHostKey(),
			}

			client, err := ssh.Dial("tcp", fmt.Sprintf("%s:%d", host.IP, host.Port), config)
			if err != nil {
				app.GetLogger().Fatal("连接失败:", err)
			}
			defer client.Close()

			session, err := client.NewSession()
			if err != nil {
				app.GetLogger().Fatal("创建会话失败:", err)
			}
			defer session.Close()

			// 使用带缓冲的输入输出处理
			session.Stdin = bufio.NewReader(os.Stdin)
			writer := bufio.NewWriter(os.Stdout)
			session.Stdout = writer
			session.Stderr = writer

			if err := session.RequestPty("xterm", 80, 24, ssh.TerminalModes{}); err != nil {
				app.GetLogger().Fatal("请求PTY失败:", err)
			}
		if err := session.Shell(); err != nil {
			app.GetLogger().Fatal("启动shell失败:", err)
		}
			session.Wait()
		},
	}
}

// 辅助函数
type Host struct {
	IP       string `json:"ip"`
	Username string `json:"username"`
	Password string `json:"password"`
	Port     int    `json:"port"`
}

func prompt(text string) string {
	fmt.Print(text)
	reader := bufio.NewReader(os.Stdin)
	input, _ := reader.ReadString('\n')
	// 转换Windows换行符为UNIX格式
	return strings.TrimRight(strings.ReplaceAll(input, "\r\n", "\n"), "\n")
}

func mustAtoi(s string) int {
	i, err := strconv.Atoi(s)
	if err != nil {
		app.GetLogger().Fatal("无效的数字输入")
	}
	return i
}

